////////////////////////////////////////////////////////////////
// multinomial expression divergence
// (c) Yongjin Park, 2013
#ifndef MED_HH_
#define MED_HH_

#include "util.hh"
#include "expression.hh"
#include "expression_tuple.hh"
#include <boost/ptr_container/ptr_vector.hpp>
#include <iostream>
#include <fstream>

// all pairwise divergence
struct
func_expr_pairwise_div_t
{
 
public:

  typedef double value_type;
  typedef std::vector<double> vec_type;
  typedef boost::ptr_vector< vec_type > mat_type;

  typedef expr_tuple_t data_t;
  typedef expr_tuple_set_t data_set_t;
  typedef expr_tuple_set_iter_t iterator_t;

  explicit
  func_expr_pairwise_div_t(const size_t p, double bmin = -5., double bmax = 5.) :
    n(0.), num_vectors(p), num_pairs(p*(p-1)/2), penalty(1e-3),
    beta_min(bmin), beta_max(bmax) {

    TLOG( "func_expr_pairwise_div_t() num_vectors = " << num_vectors << " num_pairs = " << num_pairs );

    assert_msg(p > 1, "must have at least a pair of vectors");
    assert_msg(bmin < bmax, "must have beta_min < beta_max");

    init_mat( beta_pairs, 2 );
    init_mat( beta_sq_pairs, 2 );
    init_mat( invtausq_pairs, 2, 1. );
    
    init_mat( prec_pairs, 3 );
    init_mat( cov_pairs, 3 );

    init_mat( eta_stoch_pairs, 2 );
    init_mat( prec_stoch_pairs, 3 );
  }

  virtual ~func_expr_pairwise_div_t () {}

  // non-conjugate & stochastic variational update 
  double update(const data_set_t& data_set,
		const double rate = 1.,
		const double size_factor = 1.)
    {
      const size_t max_iter = 100;
      const double tol = 1e-3;
      
      if( data_set.size() == 0 ) return 0.; // no data at all

      // non-conjugate update based on current observation
      for(size_t iter=0; iter<max_iter; ++iter)
      {
	double diff = 0.;
	for(size_t k=0; k<num_vectors; ++k)
	  for(size_t l=(k+1); l<num_vectors; ++l)
	  {
	    // optimize beta for all pairs
	    optimize_beta( data_set, k, l );

	    // update inv tau sq terms
	    diff += update_inv_tau_sq(k,l);
	  }
	if( diff < tol ) break;
      }

      // stochastic update
      // prec_stoch and eta_stoch
      // current eta = prec * beta
#ifdef DEBUG
      assert_msg( (rate > 0) & (rate <= 1.), "rate out of bound" );
      assert_msg( size_factor > 0., "invalid size" );
#endif
      double delta = 0.;

      n = (1.-rate) * n +
	rate * size_factor * ((double)data_set.size());

      for(size_t k=0; k<num_vectors; ++k)
	for(size_t l=(k+1); l<num_vectors; ++l)
	{
	  const vec_type& beta = get_vec_const( beta_pairs, k, l );
	  const vec_type& prec = get_vec_const( prec_pairs, k, l );
	  vec_type& cov = get_vec( cov_pairs, k, l );
	  
	  vec_type& prec_stoch = get_vec( prec_stoch_pairs, k, l );
	  vec_type& eta_stoch = get_vec( eta_stoch_pairs, k, l );

	  double eta0 = beta[0] * prec[pos_00] + beta[1] * prec[pos_01];
	  double eta1 = beta[0] * prec[pos_01] + beta[1] * prec[pos_11];
  
	  double eta0_old = eta_stoch[0];
	  double eta1_old = eta_stoch[1];

	  eta_stoch[0] = eta0_old*(1.-rate) + eta0 * rate * size_factor;
	  eta_stoch[1] = eta1_old*(1.-rate) + eta1 * rate * size_factor;

	  delta += std::sqrt( std::pow(eta_stoch[0] - eta0_old, 2.) +
			      std::pow(eta_stoch[1] - eta1_old, 2.) );

	  prec_stoch[pos_00] = prec_stoch[pos_00] * (1. - rate) + prec[pos_00] * rate * size_factor;
	  prec_stoch[pos_01] = prec_stoch[pos_01] * (1. - rate) + prec[pos_01] * rate * size_factor;
	  prec_stoch[pos_11] = prec_stoch[pos_11] * (1. - rate) + prec[pos_11] * rate * size_factor;

	  // re-calculate covariance matrix
	  matrix_inverse( prec_stoch, cov );
	}      

      return delta;
    }

  // simple setter and getter
  double get_penalty() const { return penalty; }
  void set_penalty(double _p){ penalty = _p; }

  size_t get_num_pairs() const { return num_pairs; }
  size_t get_num_vectors() const { return num_vectors; }

  double get_beta_abs_sum() const
    {
      double ret = 0.;
      for(size_t k=0; k<num_vectors; ++k)
	for(size_t l=(k+1); l<num_vectors; ++l)
	{
	  const vec_type& bsq = get_vec_const( beta_sq_pairs, k, l );
	  ret += std::sqrt( bsq.at(0) ) + std::sqrt( bsq.at(1) );
	}
      return ret;
    }


  // locally collapsed variational inference
  double score(const data_t& data) const
    {
      if( n < 1. ) return score_prior( data );

      double ret = 0.;
      for(size_t k=0; k<num_vectors; ++k)
	for(size_t l=(k+1); l<num_vectors; ++l)
	{
	  ret += score(data, k, l);
	}
      
      return ret;
    }

private:

  // prior score without any observation
  double score_prior( const data_t& data ) const
    {
      double ret = 0.;
      for(size_t k=0; k<num_vectors; ++k)
	for(size_t l=0; l<num_vectors; ++l)
	{
	  const vec_type& X = data.vectors.at(k);
	  const vec_type& Y = data.vectors.at(l);
	  
	  const double nA = X.size();
	  const double nB = Y.size();
	  ret += nA * nB * log_sigmoid(0.);
	}
      return ret;
    }

protected:
  
  double score(const data_t& data, const size_t k, const size_t l)
    const
    {
#ifdef DEBUG
      assert_msg( k < l, "must have ordered (k,l)" );
#endif
      const vec_type& X = data.vectors.at(k);
      const vec_type& Y = data.vectors.at(l);
      
      const double nA = X.size();
      const double nB = Y.size();

      const vec_type& prec = get_vec_const( prec_stoch_pairs, k, l );
      const vec_type& cov = get_vec_const( cov_pairs, k, l );
      const vec_type& eta = get_vec_const( eta_stoch_pairs, k, l );

      double b0 = cov[pos_00] * eta.at(0) + cov[pos_01] * eta.at(1);
      double b1 = cov[pos_01] * eta.at(0) + cov[pos_11] * eta.at(1);

// #ifdef DEBUG
//       std::cerr << "beta0 = " << b0 << ", beta1 = " << b1 << std::endl;
// #endif
      
      // temporary storage to calculate
      // M = prec + [ m00, m01 ]
      //            [ m01, m11 ]
      // where
      // m00 = nB sum wa + nA sum vb
      // m01 = nB sum wa xa + nA sum vb yb
      // m11 = nB sum wa xa^2 + nB sum vb yb^2
      //
      // N = eta + [ nB sum wa ra - nA sum vb sb ] 
      //          [ nB sum wa xa ra - nA sum vb yb sb ]

      // vec_type M(3,0.);
      // vec_type M_inv(3,0.);
      // vec_type N(2,0.);

      double M00 = 0., M01 = 0., M11 = 0.;
      double N0 = 0., N1 = 0.;

      double rWr = 0.;
      double sVs = 0.;

      M00 = prec.at(pos_00); M01 = prec.at(pos_01); M11 = prec.at(pos_11);

      for(size_t a=0; a<X.size(); ++a)
      {
	double xa = X.at(a);
	double wa = sigvar( b0 + b1 * xa );
	double ra = b0 + b1 * xa + 1. + std::exp( -b0 - b1 * xa );

	M00 += nB * wa;
	M01 += nB * wa * xa;
	M11 += nB * wa * xa * xa;

	N0 += nB * wa * ra;
	N1 += nB * wa * xa * ra;

	rWr += wa * ra* ra;
      }

      for(size_t b=0; b<Y.size(); ++b)
      {
	double yb = Y.at(b);
	double vb = sigvar( b0 + b1 * yb );
	double sb = -b0 - b1 * yb + 1. + std::exp( b0 + b1 * yb );

	M00 += nA * vb;
	M01 += nA * vb * yb;
	M11 += nA * vb * yb * yb;

	N0 -= nA * vb * sb;
	N1 -= nA * vb * yb * sb;

	sVs += vb * sb * sb;
      }

      // invert M matrix
      double iM00, iM11, iM01;
      matrix_inverse( M00, M11, M01, iM00, iM11, iM01 );

      double ret = 0.5 * N0*N0*iM00 + N0*N1*iM01 + 0.5 * N1*N1*iM11;

      ret += - nB * 0.5 * rWr - nA * 0.5 * sVs;

      ret += - 0.5 * ( eta.at(0) * b0 + eta.at(1) * b1 );
      
      ret += 0.5 * std::log( matrix_determ( prec ) )
	- 0.5 * std::log( matrix_determ( M00, M11, M01 ) );

      return ret;
    }

  // optimize beta coefficients
  // also update local precision matrix
  void optimize_beta( const data_set_t& data_set,
		      const size_t k, const size_t l )
    {

      const size_t max_iter = 1000;
      const double tol = 1e-3;

#ifdef DEBUG
      assert_msg( k < num_vectors && l < num_vectors, "index out of bound" );
      assert_msg( k < l, "must have ordered (k,l)" );
#endif

      vec_type& beta = get_vec( beta_pairs, k, l );
      vec_type& prec = get_vec( prec_pairs, k, l );
      const vec_type& invtausq = get_vec_const( invtausq_pairs, k, l );

      double b0 = beta.at(0), b1 = beta.at(1);  // old beta coefficients
      
      for(size_t iter=0; iter<max_iter; ++iter)
      {
	double num0 = 0.;  // numerator for b0
	double num1 = 0.;  // numerator for b1
	double denom0 = invtausq.at(0); // denominator for b0
	double denom1 = invtausq.at(1); // denominator for b1

	double prec_00 = invtausq.at(0); // precision
	double prec_01 = 0.;             // matrix
	double prec_11 = invtausq.at(1); //

	for(iterator_t it = begin(data_set); it != end(data_set); ++it)
	{
	  double z = it.get_prob();
	  const data_t& data = *it;
#ifdef DEBUG
	  assert_msg( data.size() == num_vectors, "incomplete data" );
#endif
	  const vec_type& X = data.vectors.at(k);
	  const vec_type& Y = data.vectors.at(l);
	
	  // const double nA = X.size();
	  // const double nB = Y.size();

	  if( z < 1e-10 ) continue;

	  for(size_t a=0; a<X.size(); ++a)
	  {
	    double xa = X.at(a);
	    double wa = sigvar( b0 + b1 * xa );
	    double ra = b0 + b1 * xa + 1. + std::exp( -b0 - b1 * xa );
	    denom0 += wa * z; // * nB
	    num0 += z * wa * (ra - b1 * xa); // nB * 
	    denom1 += z * wa * xa * xa; // nB * 
	    num1 += z * wa * (ra - b0) * xa; // nB * 

	    prec_00 += wa * z; // nB * 
	    prec_01 += wa * z * xa; // nB * 
	    prec_11 += wa * z * xa * xa; // nB * 
	  }

	  for(size_t b=0; b<Y.size(); ++b)
	  {
	    double yb = Y.at(b);
	    double vb = sigvar( b0 + b1 * yb );
	    double sb = -b0 - b1 * yb + 1. + std::exp( b0 + b1 * yb );
	    denom0 += vb * z; // * nA
	    num0 -= z * vb * (sb + b1 * yb); // nA * 
	    denom1 += z * vb * yb * yb; // nA * 
	    num1 -= z * vb * (sb + b0) * yb; // nA * 

	    prec_00 += vb * z; // nA * 
	    prec_01 += vb * z * yb; // nA * 
	    prec_11 += vb * z * yb * yb; // nA * 
	  }
	}

	b0 = num0 / std::max(1e-10, denom0);
	b1 = num1 / std::max(1e-10, denom1);

	if( b0 < beta_min ) b0 = beta_min;
	if( b0 > beta_max ) b0 = beta_max;
	if( b1 < beta_min ) b1 = beta_min;
	if( b1 > beta_max ) b1 = beta_max;

	double delta = std::sqrt( std::pow(b0-beta.at(0),2.) +
				  std::pow(b1-beta.at(1),2.) );

	beta[0] = b0;
	beta[1] = b1;

	prec[pos_00] = prec_00;
	prec[pos_01] = prec_01;
	prec[pos_11] = prec_11;

      }

      vec_type& cov = get_vec( cov_pairs, k, l );
      matrix_inverse( prec, cov );
      vec_type& beta_sq = get_vec( beta_sq_pairs, k, l );
      beta_sq[0] = cov[pos_00] + beta[0] * beta[0];
      beta_sq[1] = cov[pos_11] + beta[1] * beta[1];

// #ifdef DEBUG
//       std::cerr << k << ", " << l << "\tbeta:";
//       dump_vec( beta );
//       std::cerr << "\tprec:";
//       dump_vec( prec );
//       std::cerr << "\tcov:";
//       dump_vec( cov );
// #endif
    }

  // update inv tau sq
  double update_inv_tau_sq (const size_t k, const size_t l) 
    {
#ifdef DEBUG
      assert_msg( k != l, "no self loop" );
      assert_msg( k < num_vectors && l < num_vectors, "index out of bound" );
      assert_msg( k < l, "must have ordered (k,l)" );
#endif

      vec_type& invtausq = get_vec(invtausq_pairs, k, l);
      const vec_type& beta_sq = get_vec_const(beta_sq_pairs, k, l);

      double it0, it1;

      // inv_tau_sq = sqrt( penalty^2 / E[beta^2] )
      it0 = penalty / std::sqrt( beta_sq.at(0) );
      it1 = penalty / std::sqrt( beta_sq.at(1) );

      double ret = std::sqrt( std::pow(it0 - invtausq.at(0), 2.) +
			      std::pow(it1 - invtausq.at(1), 2.) );
      invtausq[0] = it0;
      invtausq[1] = it1;
      return ret;
    }

protected:

  double n;

  // dimensionality

  const size_t num_vectors;
  const size_t num_pairs;

  const double beta_min;
  const double beta_max;

  // variational parameters

  mat_type beta_pairs;     // k x l pairwise, each contains 1 x 2 vector
  mat_type beta_sq_pairs;  // k x l pairwise, each contains 1 x 2 vector
  mat_type invtausq_pairs; // k x l pairwise, each contains 1 x 2 vector

  mat_type prec_pairs;     // k x l pairwise, each contains 1 x 3 vector
  mat_type cov_pairs;      // k x l pairwise, each contains 1 x 3 vector

  // stochastic variational update

  mat_type eta_stoch_pairs;  // k x l pairwise, each contains 1 x 2 vector
  mat_type prec_stoch_pairs;// k x l pairwise, each contains 1 x 3 vector

  double penalty;    // L1 penalty

  // unique id for a pair (k,l) without direction
  size_t pair_id( const size_t k, const size_t l ) const
    {
#ifdef DEBUG
      assert_msg( k != l, "no self loop" );
      assert_msg( k < num_vectors, "index out of bound" );
      assert_msg( l < num_vectors, "index out of bound" );
#endif
      size_t r = k < l ? k : l;
      size_t c = k > l ? k : l;

      // unique id = sum_{j=1}^{r} (n-j-1) + c - 1
      //           = r*(n-1) - r*(r+1)/2 + c - 1
      size_t id = r * (num_vectors-1) - r*(r+1)/2 + c - 1;

      // TLOG( "pair id = " << id << " from " << k << ", " << l );
      // can verify r = (n-1), c = n
      // tot pairs = sum_{j=1}^{n-1} (n-j-1) + n - 1
      //           = n * (n-1) / 2
      return id;
    }

  vec_type& get_vec( mat_type& mat, size_t k, size_t l )
    {
      return mat.at( pair_id(k,l) );
    }

  const vec_type& get_vec_const( const mat_type& mat, size_t k, size_t l ) const
    {
      return mat.at( pair_id(k,l) );
    }

  void init_mat( mat_type& mat, const size_t ncols ){
    init_mat( mat, ncols, 0. );
  }

  void init_mat( mat_type& mat, const size_t ncols, double val ){
    for(size_t j=0; j<num_pairs; ++j)
      mat.push_back( new vec_type(ncols, val) );
  }

  // math helpers
  
  double sigmoid(const double z) const
    {
      double ret = 1. / std::min(1e4, (1. + std::exp(-z))); // [1e-4, 1]
#ifdef DEBUG
      assert_msg( !std::isnan(ret), "isnan" );
      assert_msg( !std::isinf(ret), "isnif" );
#endif
      return ret;
    }

  double log_sigmoid(const double z) const
    {
      return -std::log(1. + std::exp(-z));
    }

  double sigvar(const double z) const
    {
      return sigmoid(z) * sigmoid(-z);
    }

  // matrix helper

  static const size_t pos_00 = 0;
  static const size_t pos_01 = 1;
  static const size_t pos_11 = 2;

  static void set_matrix( const vec_type& src, vec_type& tgt )
    {
      tgt[ pos_00 ] = src.at( pos_00 );
      tgt[ pos_01 ] = src.at( pos_01 );
      tgt[ pos_11 ] = src.at( pos_11 );
    }

  static double matrix_determ( const vec_type& m )
    {
      double a = m.at(pos_00), b = m.at(pos_11), c = m.at(pos_01);
      double det = std::abs(a*b - c*c);
      return det;
    }

  static double matrix_determ( const double a, const double b, const double c )
    {
      double det = std::abs(a*b - c*c);
      return det;
    }

  // 2 x 2 symmetric matrix inverse
  // | a c | ^-1 = | b -c | * 1 / (a*b - c*c)
  // | c b |       | -c a |  
  static void matrix_inverse( const vec_type& org, vec_type& inv )
    {
#ifdef DEBUG
      assert_msg( org.size() == inv.size(), "size should match" );
#endif
      double a = org[pos_00], b = org[pos_11], c = org[pos_01];
      double det = std::abs(a*b - c*c);
      inv[pos_00] = b / det;
      inv[pos_01] = -c / det;
      inv[pos_11] = a / det;
    }

  static void matrix_inverse( const double a, const double b, const double c,
			      double& ainv, double& binv, double& cinv )
    {
      double det = std::abs( a*b - c*c );
      ainv = b / det;
      cinv = -c / det;
      binv = a / det;
    }

public:
  // output for each line
  // k l E[beta0] E[beta1] sd(beta0) sd(beta1)

  // void
  // write_beta( std::ostream & ofs )	      
  //   const
  //   {
  //     for(size_t k=0; k<num_vectors; ++k)
  // 	for(size_t l=(k+1); l<num_vectors; ++l)
  // 	{
  // 	  const vec_type& mean = get_vec_const( beta_pairs, k, l );
  // 	  const vec_type& cov = get_vec_const( cov_pairs, k, l );
  // 	  ofs << k << "\t" << l << "\t"
  // 	      << mean.at(0) << "\t" << std::sqrt(cov.at( pos_00 ))
  // 	      << "\t"
  // 	      << mean.at(1) << "\t" << std::sqrt(cov.at( pos_11 ))
  // 	      << std::endl;
  // 	}
  //   }

  void
  write_beta( std::ostream & ofs )	      
    const
    {
      for(size_t k=0; k<num_vectors; ++k)
	for(size_t l=(k+1); l<num_vectors; ++l)
	{
	  const vec_type& cov = get_vec_const( cov_pairs, k, l );
	  const vec_type& eta = get_vec_const( eta_stoch_pairs, k, l );
	  
	  double b0 = cov[pos_00] * eta.at(0) + cov[pos_01] * eta.at(1);
	  double b1 = cov[pos_01] * eta.at(0) + cov[pos_11] * eta.at(1);
	  	  
	  ofs << k << "\t" << l << "\t"
	      << b0 << "\t" << std::sqrt(cov.at( pos_00 )) << "\t"
	      << b1 << "\t" << std::sqrt(cov.at( pos_11 )) << std::endl;
	}
    }

};

// empirical Bayes estimate of penalty parameter
double empirical_bayes_med( func_expr_pairwise_div_t& obj,
			    double rate = 1. )
{
  double beta_abs_sum = obj.get_beta_abs_sum();
  double tol = 1e-3;
  size_t max_iter = 1000;
  
  double penalty = obj.get_penalty();
  double n = obj.get_num_pairs();
  if( penalty < 1e-10 )
  {
    penalty = n / beta_abs_sum;
  }
  else
  {
    for(size_t iter = 0; iter < max_iter; ++iter)
    {
      double penalty_old = penalty;
      double stuff = 0.5 / penalty * penalty + 0.5 * beta_abs_sum / n;
      penalty = 1 / std::sqrt( stuff );
      double diff = std::sqrt(std::pow(penalty-penalty_old,2.));
      diff /= std::max( 1e-10, penalty_old );
      if( diff < tol ) break;
    }
  }

  penalty = (1. - rate) * obj.get_penalty() + rate * penalty;
  double diff = std::sqrt( std::pow( penalty - obj.get_penalty(), 2. ) );
  obj.set_penalty( penalty );
  return diff;
}

double empirical_bayes( func_expr_pairwise_div_t& obj,
			double rate = 1. )
{
  return empirical_bayes_med( obj, rate );
}
#endif
